from math import exp

from GA.FitnessFunction import FitnessFunction


class accuracy(FitnessFunction):
    def __init__(self):
        pass

    def __call__(self, data, predictions, label_key="label", pos_label='Y'):
        """计算准确率，召回率，F1值"""
        pred_labels, _ = predictions
        assert len(data) == len(pred_labels)
        correct = 0
        total = 0
        tp = 0
        fp = 0
        fn = 0
        for i in range(len(data)):
            total += 1
            if data[i][label_key] == pred_labels[i]:
                correct += 1
                if data[i][label_key] == pos_label:
                    tp += 1
            else:
                if pred_labels[i] == pos_label:
                    fp += 1
                else:
                    fn += 1
        if tp == 0:
            precision = 0
            recall = 0
            f1 = 0
        else:
            precision = tp / (tp + fp)
            recall = tp / (tp + fn)
            f1 = (2 * precision * recall) / (precision + recall)
        acc = correct / total
        extra = {
            "acc": acc,
            "precision": precision,
            "recall": recall,
            "f1": f1
        }
        return acc, extra


class mean_hinge_probs(FitnessFunction):
    def __init__(self):
        pass

    def __call__(self, data, predictions, label_key="label"):
        """计算平均 hinge loss"""
        pred_labels, pred_logprobs = predictions
        assert len(data) == len(pred_labels)
        total = 0
        for i in range(len(data)):
            if data[i][label_key] == pred_labels[i]:
                total += exp(pred_logprobs[i])
            else:
                total += 0
        score = total / len(data)
        return score, {"score": score}
